#include <iostream>
#include <fstream>
#include <string>
#include <map>
#include <dirent.h>
#include <sys/types.h>
#include <signal.h>
#include "../include/proc_utils.h"
using namespace std;

ProcNode::ProcNode(int pid): m_pid(pid), m_children(vector<ProcNode*>()) {}

int ProcNode::get_ppid() {
    ifstream fin;
    string path = "/proc/" + to_string(m_pid) + "/status";
    fin.open(path, ios::in);
    if (!fin.is_open()) {
        cout << "打开文件错误" << endl;
        return 0;
    }
    string buff;
    while (getline(fin, buff)) {
        if (buff.substr(0, 4).compare("PPid") == 0) {
            break;
        }
    }
    fin.close();
    return stoi(buff.substr(6));
}

ProcTree::ProcTree() {
    map<int, ProcNode*> proc_map; 

    DIR *dirptr = NULL;
    struct dirent *entry = NULL;
    dirptr = opendir("/proc");
    int count = 0;
    while (entry = readdir(dirptr)) {
        int pid;
        if (entry->d_type == DT_DIR && (pid = atoi(entry->d_name))) {
            proc_map[pid] = new ProcNode(pid);
        }
    }
    closedir(dirptr);
    for (map<int, ProcNode*>::iterator it = proc_map.begin(); it != proc_map.end(); it++) {
        int ppid = it->second->get_ppid();
        if (ppid != 0) 
            proc_map[ppid]->m_children.emplace_back(it->second);
    }
    root = proc_map[1];
}

vector<int> ProcTree::get_descendants_pids(int ppid) {
    // find the process whose id is equal to ppid;
    vector<ProcNode*> stack;
    ProcNode* parent = NULL;
    ProcNode* node = NULL;
    if (root->m_pid == ppid) {
        parent = root;
    } else {
        stack.emplace_back(root);
        while (!stack.empty()) {
            node = stack.back();
            stack.pop_back();
            if (node->m_pid == ppid) {
                parent = node;
                break;
            } else {
                for (auto it = node->m_children.begin(); it != node->m_children.end(); it++)
                    stack.emplace_back(*it);
            }
        }
    }
    if (parent == NULL) return vector<int>();
    vector<int> result;
    stack.clear();
    stack.emplace_back(parent);
    while (!stack.empty()) {
        node = stack.back();
        stack.pop_back();
        result.emplace_back(node->m_pid);
        for (auto it = node->m_children.begin(); it != node->m_children.end(); it++)
            stack.emplace_back(*it);
    }
    return result;
}

bool ProcTree::kill_descendants(int ppid) {
    vector<int> pids = get_descendants_pids(ppid);
    if (!pids.empty()) {
        for (auto it = pids.begin(); it != pids.end(); it++) {
            if (kill(*it, SIGKILL) != 0)
                return false;
        }
    }
    return true;
}

void error_handling(char *message)
{
    fputs(message, stderr);
    fputc('\n', stderr);
    exit(1);
}